# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/14 17:31 
# @Author  : shutao
# @FileName: facenet.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import paddle
import paddle.nn as nn

from paddle import ParamAttr
from paddle.regularizer import L2Decay

from ..base_net import BaseNet
from ..modules import Conv2D


class FaceNet(BaseNet):
    def __init__(self):
        super(FaceNet, self).__init__()

        self.conv1 = Conv2D(3, 64, 7, 2, 3)
        self.pool1 = nn.MaxPool2D(3, 2, 1)
        self.rnorm1 = nn.BatchNorm2D(64)

        self.conv2a = Conv2D(64, 64, 1)
        self.conv2 = Conv2D(64, 192, 3, 1, 1)
        self.rnorm2 = nn.BatchNorm2D(192)
        self.pool2 = nn.MaxPool2D(3, 2, 1)

        self.conv3a = Conv2D(192, 192, 1)
        self.conv3 = Conv2D(192, 384, 3, 1, 1)
        self.pool3 = nn.MaxPool2D(3, 2, 1)

        self.conv4a = Conv2D(384, 384, 1)
        self.conv4 = Conv2D(384, 256, 3, 1, 1)

        self.conv5a = Conv2D(256, 256, 1)
        self.conv5 = Conv2D(256, 256, 3, 1, 1)

        self.conv6a = Conv2D(256, 256, 1)
        self.conv6 = Conv2D(256, 256, 3, 1, 1)
        self.pool4 = nn.MaxPool2D(3, 2, 1)

        self.fc1 = nn.Linear(256*7*7, 128*32)
        self.fc2 = nn.Linear(128*32, 128*32)
        self.fc7128 = nn.Linear(128*32, 128*1)
        self.L2_norm = nn.BatchNorm1D(128, weight_attr=ParamAttr(regularizer=L2Decay(coeff=0.01)))
        pass

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.pool1(x)
        x = self.rnorm1(x)

        x = self.conv2a(x)
        x = self.conv2(x)
        x = self.rnorm2(x)
        x = self.pool2(x)

        x = self.conv3a(x)
        x = self.conv3(x)
        x = self.pool3(x)

        x = self.conv4a(x)
        x = self.conv4(x)

        x = self.conv5a(x)
        x = self.conv5(x)

        x = self.conv6a(x)
        x = self.conv6(x)
        x = self.pool4(x)

        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc7128(x)
        outputs = self.L2_norm(x)
        return outputs
